//
// Created by jack on 2021/2/4.
//

#include "weilapb_defines.h"

#include <aes.h>
#include <base64.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include "weilapb_cryptutil.h"

#include "weilapb_internal.h"

PB_CHAR gKey[32];
PB_BOOL keyGenerated = PB_FALSE;

PB_UINT32 weilapb_readUint32(const PB_UINT8* buf) {
    PB_UINT32 value = (buf[0] << 24) | (buf[1] << 16) | (buf[2] << 8) | buf[3];
    return value;
}

void weilapb_writeUint32(PB_UINT8* buf, PB_UINT32 value) {
    buf[0] = (PB_UINT8)(value >> 24);
    buf[1] = (PB_UINT8)((value >> 16) & 0xFF);
    buf[2] = (PB_UINT8)((value >> 8) & 0xFF);
    buf[3] = (PB_UINT8)(value & 0xFF);
}

const PB_CHAR* weilapb_getKey() {
    if (keyGenerated == PB_FALSE) {
        int i, j, c;
        c = 0;
        for (i = 0; i < 3; i++) {
            for (j = 0; j < 10; j++) {
                if (j != 9) {
                    gKey[c++] = (PB_CHAR)('1' + j);
                } else {
                    gKey[c++] = '0';
                }
            }
        }

        gKey[c++] = '1';
        gKey[c++] = '2';

        keyGenerated = PB_TRUE;
    }

    return gKey;
}

PB_INT32 weilapb_encryptText_nomem(PB_CHAR *pText, PB_CHAR *pOutText, PB_UINT32 maxBufferSize) {
    AES_KEY        aesKey;
    const PB_CHAR* key = weilapb_getKey();
    PB_UINT8       encData[16];
    PB_UINT32      i, j;
    PB_UINT32      inputSize   = 0;
    PB_UINT32      nRemain     = 0;
    PB_UINT32      nBlocks     = 0;
    PB_UINT32      nEncryptLen = 0;
    PB_CHAR        *pEncryptedBuf = NULL;

    if (pText == NULL || pOutText == NULL) {
        return -1;
    }

    inputSize   = strlen(pText);
    nRemain     = inputSize % 16;
    nBlocks     = inputSize / 16;

    if (nRemain > 0) {
        nBlocks += 1;
    }

    nEncryptLen = nBlocks * 16;
    if (nEncryptLen * 2 + 2 > maxBufferSize) {
        return -2;
    }

    memset(pOutText, 0, maxBufferSize);
    pEncryptedBuf = pOutText + (nEncryptLen + 1) * 2;
    memcpy(pEncryptedBuf, pText, inputSize);

    AES_set_encrypt_keys((const PB_UINT8*)key, 256, &aesKey);
    for (i = 0; i < nBlocks; i++) {
        char *pTemp = pOutText + i * 32;
        memset(encData, 0, 16);
        AES_encrypts((pEncryptedBuf + i * 16), encData, &aesKey);
        for(j = 0; j < 16; j++) {
            sprintf(pTemp, "%02x", *(pEncryptedBuf + i * 16 + j));
            pTemp += 2;
        }
    }

    pOutText[nEncryptLen * 2] = 0;

    return nEncryptLen * 2;
}

PB_INT32
weilapb_encryptText(PB_CHAR* pText, PB_UINT32 inputSize, PB_CHAR** pEncryptedText, PB_UINT32* pEncryptedTextSize) {
    const PB_CHAR* key = keyGenerated ? gKey : weilapb_getKey();
    PB_UINT8       encData[16];
    PB_INT32       ouputBufSize;
    AES_KEY        aesKey;
    PB_UINT8*      pData;
    PB_UINT32      i;
    PB_UINT32      nRemain     = inputSize % 16;
    PB_UINT32      nBlocks     = (inputSize + 15) / 16;
    PB_UINT32      nEncryptLen = 0;

    if (pText == NULL || inputSize <= 0) {
        return -1;
    }

    if (nRemain > 12 || nRemain == 0) {
        nBlocks += 1;
    }

    nEncryptLen = nBlocks * 16;

    pData = (PB_UINT8*)PB_ALLOC(nEncryptLen);
    if (pData == PB_NULL) {
        return -2;
    }
    memset(pData, 0, nEncryptLen);
    memcpy(pData, pText, inputSize);

    weilapb_writeUint32((pData + nEncryptLen - 4), inputSize);
    AES_set_encrypt_keys((const PB_UINT8*)key, 256, &aesKey);
    for (i = 0; i < nBlocks; i++) {
        memset(encData, 0, 16);
        AES_encrypts(pData + i * 16, encData, &aesKey);
        memcpy(pData + i * 16, encData, 16);
    }

    ouputBufSize    = wl_base64_encode_output_size(nEncryptLen) + 16;
    *pEncryptedText = (PB_CHAR*)PB_ALLOC(ouputBufSize);
    if (*pEncryptedText == PB_NULL) {
        PB_FREE(pData);
        return -3;
    }
    memset(*pEncryptedText, 0, ouputBufSize);
    *pEncryptedTextSize = ouputBufSize;

    if (wl_base64_encode((const PB_CHAR*)pData, nEncryptLen, *pEncryptedText, pEncryptedTextSize) == 0) {
        PB_FREE(pData);
        return 0;
    }

    PB_FREE(*pEncryptedText);
    PB_FREE(pData);

    return -4;
}

PB_INT32 weilapb_decryptText(PB_CHAR*   pCryptedText,
                             PB_UINT32  inputSize,
                             PB_CHAR**  pDecryptedText,
                             PB_UINT32* pDecryptedTextSize) {
    const char* key   = keyGenerated ? gKey : weilapb_getKey();
    PB_UINT8*   pData = PB_NULL;
    PB_UINT32   outputBufferSize;

    if (pCryptedText == PB_NULL || inputSize <= 0) {
        return -1;
    }

    pData            = (PB_UINT8*)PB_ALLOC(inputSize);
    outputBufferSize = inputSize;
    if (wl_base64_decode(pCryptedText, strlen(pCryptedText), (PB_CHAR*)pData, &outputBufferSize) == 0) {
        PB_UINT32 i;
        AES_KEY   aesKey;
        PB_UINT32 nBlocks;
        char      decodedBlock[16];
        if (outputBufferSize == 0) {
            PB_FREE(pData);
            return -2;
        }

        if (outputBufferSize % 16 != 0) {
            PB_FREE(pData);
            return -3;
        }

        nBlocks = outputBufferSize / 16;
        AES_set_decrypt_keys((const PB_UINT8*)key, 256, &aesKey);  //设置AES解密密钥
        for (i = 0; i < nBlocks; i++) {
            memset(decodedBlock, 0, 16);
            AES_decrypts((const PB_UINT8*)(pData + i * 16), (PB_UINT8*)decodedBlock, &aesKey);
            memcpy(pData + i * 16, decodedBlock, 16);
        }

        *pDecryptedTextSize = weilapb_readUint32((pData + outputBufferSize - 4));
        memset(pData + (*pDecryptedTextSize), 0, inputSize - (*pDecryptedTextSize));
        *pDecryptedText = (PB_CHAR*)pData;

        return 0;
    }

    PB_FREE(pData);
    return -1;
}
